import pandas as pd
from ast import literal_eval
import json

# 读取原始数据
df = pd.read_csv('tourism2.csv', converters={
    'Interests': literal_eval,
    'Sites Visited': literal_eval
})

# 计算平均满意度
A = df['Satisfaction'].mean()

# 生成评分数据
rating_df = (
    df.explode('Sites Visited')
    .assign(Rating=lambda x: A + 0.25)
    [['Tourist ID', 'Age', 'Interests', 'Preferred Tour Duration', 'Sites Visited', 'Rating']]
    .drop_duplicates(['Tourist ID', 'Sites Visited'])
    .reset_index(drop=True)
)


# 数据存储函数
def save_travel_data(data, filename, format='csv'):
    """
    保存旅游数据到指定格式的文件
    参数：
        data: 要保存的DataFrame
        filename: 文件名（不带扩展名）
        format: 存储格式 ['csv', 'excel', 'json']
    """
    # 转换兴趣列为字符串表示
    formatted_data = data.copy()
    formatted_data['Interests'] = formatted_data['Interests'].apply(json.dumps)

    try:
        if format == 'csv':
            formatted_data.to_csv(f"{filename}.csv", index=False)
        elif format == 'excel':
            formatted_data.to_excel(f"{filename}.xlsx", index=False, engine='openpyxl')
        elif format == 'json':
            # 保存为面向记录的JSON格式
            formatted_data.to_json(f"{filename}.json", orient='records', indent=2)
        else:
            raise ValueError("不支持的存储格式，请选择 csv/excel/json")
        print(f"数据已成功保存到 {filename}.{format}")
    except Exception as e:
        print(f"存储失败：{str(e)}")


# 保存数据（示例）
save_travel_data(rating_df, 'travel_ratings', 'csv')  # CSV格式
# save_travel_data(rating_df, 'travel_ratings', 'excel')  # Excel格式
# save_travel_data(rating_df, 'travel_ratings', 'json')  # JSON格式


# 验证数据加载
def load_travel_data(filename, format='csv'):
    """
    加载存储的旅游数据
    返回包含原始数据类型的DataFrame
    """
    try:
        if format == 'csv':
            df = pd.read_csv(f"{filename}.csv", converters={
                'Interests': lambda x: json.loads(x.replace("'", "\""))
            })
        elif format == 'excel':
            df = pd.read_excel(f"{filename}.xlsx", engine='openpyxl', converters={
                'Interests': lambda x: json.loads(x.replace("'", "\""))
            })
        elif format == 'json':
            df = pd.read_json(f"{filename}.json", orient='records')
            df['Interests'] = df['Interests'].apply(json.loads)
        return df
    except Exception as e:
        print(f"加载失败：{str(e)}")
        return None


# 验证数据加载
loaded_data = load_travel_data('travel_ratings', 'csv')
print("\n加载的验证数据：")
print(loaded_data.head())